Skip to content

Update StaticInputs coordinate handling#998

Merged
frodre merged 4 commits intomainfrom
refactor/static-input-coordinates
Mar 23, 2026
Merged

Update StaticInputs coordinate handling#998
frodre merged 4 commits intomainfrom
refactor/static-input-coordinates

Conversation

@frodre
Copy link
Collaborator

@frodre frodre commented Mar 21, 2026

PR 2/3 for #971 split

Changes:

  • Move lat/lon coordinates from individual StaticInput fields to the parent StaticInputs container

  • StaticInput.subset now takes slice arguments directly instead of ClosedIntervals

  • Extract _load_coords_from_ds with flexible coordinate name matching (lat/latitude/grid_yt, etc.)

  • Add StaticInputs.from_state_backwards_compatible to handle legacy checkpoints (coords per-field) and config-based loading (no coords in checkpoint)

  • load_static_inputs now requires at least one field and validates coordinate consistency across fields

  • Tests added

Base automatically changed from refactor/closed-interval-subset to main March 23, 2026 21:30
frodre added 2 commits March 23, 2026 14:34
  fme/downscaling/data/static.py — major refactor:
  - StaticInput: removed coords field; now just holds data; subset() takes slices (not intervals)
  - Added _load_coords_from_ds(), load_fine_coords_from_path() for loading coords from files
  - Added _has_legacy_coords_in_state(), _has_coords_in_state(), _sync_state_coordinates() for backwards-compat state migration
  - StaticInputs: added coords: LatLonCoordinates as first-class field; removed per-field coord validation; subset() now returns updated
  coords; to_device() moves coords too; get_state() includes top-level "coords"; from_state() raises if coords absent; added
  from_state_backwards_compatible()
  - _get_normalized_static_input(): returns (StaticInput, LatLonCoordinates) tuple
  - load_static_inputs(): no longer accepts None; validates coord consistency across fields

  fme/downscaling/data/test_static.py — full replacement with new test suite covering all new APIs

  fme/downscaling/data/__init__.py — added load_fine_coords_from_path export

  fme/downscaling/test_models.py — updated make_static_inputs() to new StaticInputs interface

  fme/downscaling/inference/test_inference.py — updated get_static_inputs() helper

  fme/downscaling/train.py — guarded load_static_inputs call against None config# Please enter the commit message for your changes. Lines starting
@frodre frodre force-pushed the refactor/static-input-coordinates branch from ff04394 to 665dfb9 Compare March 23, 2026 21:43
)

@classmethod
def from_state_backwards_compatible(
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This method gets used in #999, but there will be a CheckpointModelConfig gate on calling this without coords available in the state or the configuration.

@frodre frodre marked this pull request as ready for review March 23, 2026 22:12
Copy link
Contributor

@AnnaKwa AnnaKwa left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mostly LGTM, one request that doesn't need to be reviewed.

"coords": self.coords.get_state(),
}

def load_fine_coords_from_path(path: str) -> LatLonCoordinates:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like this code is duplicated in _get_normalized_static_input. Can this helper and its tests be removed? I think _load_coords_from_ds is simple enough that it doesn't need multiple tests.

@frodre frodre enabled auto-merge (squash) March 23, 2026 22:54
@frodre frodre merged commit 64bcaf9 into main Mar 23, 2026
7 checks passed
@frodre frodre deleted the refactor/static-input-coordinates branch March 23, 2026 23:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants